import os
from collections import Counter
import json
import base64
import numpy as np

from urllib import request, parse

def encodeImg(imgName):
    with open(imgName, 'rb') as f:
        data = f.read()
        encodestr = base64.b64encode(data)
    return encodestr

def detect_image(img, url=None, tenantId='10104'):
    if url is None:
        url = 'http://10.0.0.15:9996/detService/layerDet'
    headers = dict(tenantId=tenantId)
    data = {'base64Data': encodeImg(img)}
    data = parse.urlencode(data).encode('utf-8')
    req = request.Request(url, headers=headers, data=data)
    page = request.urlopen(req).read()
    page = bytes.decode(page)
    page = json.loads(page)
    return page['data']

def fmg_image_detect(left_img, right_img, url=None, tenantId='10104', left_border=607, right_border=319):
    left_result = detect_image(left_img, url, tenantId)
    left_result.update(dict(border=left_border))
    right_result = detect_image(right_img, url, tenantId)
    right_result.update(dict(border=right_border))
    output = dict(left=left_result, right=right_result)
    return output

def fmg_detect_concat(data):
    # data:
    #   dict which includes left and right detect results,
    #   and contains divided axis lines each side.
    left = data['left']
    right = data['right']
    left_skus = catch_skus_with_border(left, side='left')
    right_skus = catch_skus_with_border(right, side='right')
    align_right_skus_to_left(left_skus, right_skus)

    left_border_bboxes = left_skus['border_bboxes']
    left_border_scores = left_skus['border_scores']
    right_border_bboxes = right_skus['border_bboxes']
    right_border_scores = right_skus['border_scores']
    overlaps = bboxes_overlaps(left_border_bboxes, right_border_bboxes, mode='minimum')
    left_keep_ids = []
    right_keep_ids = []
    # need to refine
    iou_thr = 0.25

    for i, iou in enumerate(overlaps):
        ids = np.where(iou > iou_thr)[0]
        if not len(ids):
            left_keep_ids.append(i)
            continue
        area_left_i = calculate_bbox_area(left_border_bboxes[i])
        areas_right = calculate_bbox_area(right_border_bboxes[ids])
        j = np.argmax(areas_right)
        area_right_j = areas_right[j]
        if area_left_i > area_right_j:
            left_keep_ids.append(i)
        elif area_left_i < area_right_j:
            right_keep_ids.append(ids[j])
        else:
            score_left_i = left_border_scores[i]
            score_right_j = right_border_scores[j]
            if score_left_i >= score_right_j:
                left_keep_ids.append(i)
            else:
                right_keep_ids.append(ids[j])

    for i, iou in enumerate(overlaps.T):
        ids = np.where(iou > iou_thr)[0]
        if not len(ids):
            right_keep_ids.append(i)

    right_keep_ids = set(right_keep_ids)
    left_keep_ids = set(left_keep_ids)
    # merge boxInfo
    boxInfo = left_skus['inside_boxInfo'] + right_skus['inside_boxInfo']
    boxInfo += [left_skus['border_boxInfo'][i] for i in left_keep_ids]
    boxInfo += [right_skus['border_boxInfo'][i] for i in right_keep_ids]
    # format results
    width = left['border'] + right_skus['width'] - right['border']
    assert left_skus['height'] == left_skus['height']
    imageInfo = left['imageInfo']
    imageInfo['width'] = width
    layerNum = max(left['layerNum'], right['layerNum'])
    skuStat = calculate_skuStat(boxInfo)

    result = dict(boxInfo=boxInfo, skuStat=skuStat, imageInfo=imageInfo, layerNum=layerNum)
    return result

def calculate_skuStat(boxInfo):
    skuCount = dict(Counter([box['skuName'] for box in boxInfo]))
    skuStat = [[skuName, count] for skuName, count in skuCount.items()]
    return skuStat

def calculate_bbox_area(bbox):
    if len(bbox.shape) == 1:
        return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
    return (bbox[:,2] - bbox[:,0]) * (bbox[:,3] - bbox[:,1])

def align_right_skus_to_left(left, right):
    new_cat_width = left['border'] + right['width'] - right['border']
    new_cat_height = left['height']
    # clip left bboxes which out of boundary
    clip_left_boxInfo(left, new_cat_width, new_cat_height)

    # offset_right_result
    offset_right_result_and_clip_boxInfo(right, left, new_cat_width, new_cat_height)

def clip_left_boxInfo(left, new_cat_width, new_cat_height):
    border_ignore_ids = []
    for boxInfo in [left['border_boxInfo'], left['inside_boxInfo']]:
        new_boxInfo = []
        for i, box in enumerate(boxInfo):
            xmin, ymin, xmax, ymax = list(box['location'].values())
            w1 = xmax - xmin
            xmax = min(xmax, left['border'])
            w2 = xmax - xmin
            if w2 < 0.5 * w1:
                border_ignore_ids.append(i)
                continue
            xmin /= new_cat_width
            xmax /= new_cat_width
            ymin /= new_cat_height
            ymax /= new_cat_height
            box['location']['xmin'] = xmin
            box['location']['ymin'] = ymin
            box['location']['xmax'] = xmax
            box['location']['ymax'] = ymax
            new_boxInfo.append(box)
        boxInfo[:] = new_boxInfo

    # left['border_boxInfo'][:] = [box for i, box in enumerate(left['border_boxInfo']) if i not in border_ignore_ids]
    if border_ignore_ids:
        num_border_bboxes = len(left['border_bboxes'])
        border_keep_ids = [i for i in range(num_border_bboxes) if i not in border_ignore_ids]
        left['border_bboxes'] = left['border_bboxes'][border_keep_ids]
        left['border_scores'] = left['border_scores'][border_keep_ids]

def offset_right_result_and_clip_boxInfo(right, left, new_cat_width, new_cat_height):
    border_ignore_ids = []
    # boxInfo
    for boxInfo in [right['border_boxInfo'], right['inside_boxInfo']]:
        new_boxInfo = []
        for i , box in enumerate(boxInfo):
            xmin, ymin, xmax, ymax = list(box['location'].values())
            xmin = left['border'] + xmin - right['border']
            xmax = left['border'] + xmax - right['border']
            w1 = xmax - xmin
            xmin = max(xmin, left['border'])
            w2 = xmax - xmin
            if w2 < 0.5 * w1:
                border_ignore_ids.append(i)
                continue
            xmin /= new_cat_width
            xmax /= new_cat_width
            ymin /= new_cat_height
            ymax /= new_cat_height
            box['location']['xmin'] = xmin
            box['location']['ymin'] = ymin
            box['location']['xmax'] = xmax
            box['location']['ymax'] = ymax
            new_boxInfo.append(box)
        boxInfo[:] = new_boxInfo

    # bboxes
    right['border_bboxes'][:, 0][:] = left['border'] + right['border_bboxes'][:, 0] - right['border']
    right['border_bboxes'][:, 2][:] = left['border'] + right['border_bboxes'][:, 2] - right['border']

    # right['border_boxInfo'][:] = [box for i, box in enumerate(right['border_boxInfo']) if i not in border_ignore_ids]
    if border_ignore_ids:
        num_border_bboxes = len(right['border_bboxes'])
        border_keep_ids = [i for i in range(num_border_bboxes) if i not in border_ignore_ids]
        right['border_bboxes'] = right['border_bboxes'][border_keep_ids]
        right['border_scores'] = right['border_scores'][border_keep_ids]

def catch_skus_with_border(result, side='right'):
    assert side in ('left', 'right')
    border = result['border']
    boxInfo = result['boxInfo']
    imageInfo = result['imageInfo']
    width = imageInfo['width']
    height = imageInfo['height']
    bboxes, scores = boxInfo2bboxes(width, height, boxInfo)
    condition = (bboxes[:, 0] < border) * (bboxes[:, 2] > border)
    border_ids = np.where(condition)[0]

    if side == 'right':
        condition = bboxes[:, 0] >= border
        inside_ids = np.where(condition)[0]
    else:
        condition = bboxes[:, 2] <= border
        inside_ids = np.where(condition)[0]

    inside_boxInfo = [box for i, box in enumerate(boxInfo) if i in inside_ids]
    border_boxInfo = [box for i, box in enumerate(boxInfo) if i in border_ids]
    border_bboxes = bboxes[border_ids]
    border_scores = scores[border_ids]

    output = dict(width=width, height=height, border=border,
                  inside_ids=inside_ids, inside_boxInfo=inside_boxInfo,
                  border_ids=border_ids, border_boxInfo=border_boxInfo,
                  border_bboxes=border_bboxes, border_scores=border_scores,
                  )
    return output

def boxInfo2bboxes(width, height, boxInfo):
    bboxes, scores = [], []
    for box in boxInfo:
        location = box['location']
        xmin, ymin, xmax, ymax = list(location.values())
        xmin *= width
        xmax *= width
        ymin *= height
        ymax *= height
        location['xmin'] = int(xmin)
        location['xmax'] = int(xmax)
        location['ymin'] = int(ymin)
        location['ymax'] = int(ymax)
        bboxes.append(list(location.values()))
        scores.append(box['score'])
    bboxes = np.array(bboxes)
    scores = np.array(scores)
    return bboxes, scores

def bboxes_overlaps(bboxes1, bboxes2, mode='standard'):
    assert mode in ('standard', 'minimum')
    if isinstance(bboxes1, list):
        bboxes1 = np.array(bboxes1)
    if isinstance(bboxes2, list):
        bboxes2 = np.array(bboxes2)
    iou = np.zeros((len(bboxes2), len(bboxes1)))
    for i in range(len(bboxes2)):
        bboxes = bboxes2[i][None, :]
        bboxes = np.repeat(bboxes, len(bboxes1), axis=0)
        x1 = np.concatenate((bboxes[:, 0:1], bboxes1[:, 0:1]), axis=1).max(1)
        y1 = np.concatenate((bboxes[:, 1:2], bboxes1[:, 1:2]), axis=1).max(1)
        x2 = np.concatenate((bboxes[:, 2:3], bboxes1[:, 2:3]), axis=1).min(1)
        y2 = np.concatenate((bboxes[:, 3:4], bboxes1[:, 3:4]), axis=1).min(1)
        dw = x2 - x1
        dw[dw < 0] = 0
        dh = y2 - y1
        dh[dh < 0] = 0
        inter = dw * dh
        areas1 = (bboxes[:, 2] - bboxes[:, 0]) * (bboxes[:, 3] - bboxes[:, 1])
        areas2 = (bboxes1[:, 2] - bboxes1[:, 0]) * \
            (bboxes1[:, 3] - bboxes1[:, 1])
        if mode == 'standard':
            union = areas1 + areas2 - inter
        else:
            union = np.concatenate(
                (areas1[:, None], areas2[:, None]), 1).min(1)
        iui = inter / union
        iou[i] = iui
    return iou.T

def data_loader(src):
    imgs = [os.path.join(src, img) for img in os.listdir(src) if img.endswith('.jpg')]
    imgs.sort()
    lefts = [img for img in imgs if os.path.basename(img).startswith('L')]
    rights = [img for img in imgs if os.path.basename(img).startswith('R')]
    for left, right in zip(lefts, rights):
        yield left, right

def concate_left_right_imgs(lm, rm, cat_dir, result=None, lp=607, rp=319):
    if isinstance(lm, str):
        lm = Image.open(lm)
    if isinstance(rm, str):
        rm = Image.open(rm)
    assert lm.size[1] == rm.size[1]
    lm = np.array(lm, dtype=np.uint8)
    rm = np.array(rm, dtype=np.uint8)
    _lm = lm[:, :lp, :]
    _rm = rm[:, rp:, :]
    im = np.concatenate((_lm, _rm), 1)
    im = Image.fromarray(im)
    draw = ImageDraw.Draw(im)
    draw.line((lp,0, lp,im.size[1]), fill=128)
    if not os.path.exists(os.path.dirname(cat_dir)):
        # import shutil
        # shutil.rmtree(os.path.dirname(cat_dir))
        os.makedirs(os.path.dirname(cat_dir))
    im.save(cat_dir)
    if result is not None:
        ann_dir = os.path.splitext(cat_dir)[0] + '.json'
        with open(ann_dir, 'w') as fp:
            json.dump(result, fp)

def fmg_detect_client(data):
    url = 'http://172.16.74.199:9998/fmgLayerDet'
    url = 'http://10.0.0.15:9990/fmgLayerDet'
    # data = {k:str(v) for k, v in data.items()}
    data = parse.urlencode(data).encode('utf-8')
    req = request.Request(url, headers={'tenantId':'10104'}, data=data)
    page = request.urlopen(req).read()
    page = bytes.decode(page)
    page = json.loads(page)
    return page['data']

if __name__ == '__main__':
    from PIL import Image
    from PIL import ImageDraw
    from tqdm import tqdm

    src = 'demo/org'
    dst = 'demo/cat'
    # src = r'J:\all'
    # dst = r'J:\all\concat'

    tbar = tqdm(enumerate(data_loader(src), 1))
    L = 560
    L = 607
    R = 330
    R = 319
    for i, (left, right) in tbar:
        tbar.desc = '%d %s %s' % (i, left, right)
        result = fmg_image_detect(left, right, url=None, tenantId='10104', left_border=L, right_border=R)
        print(result)
        # result2 = fmg_detect_concat(result)
        result2 = fmg_detect_client(result)
        cat_dir = os.path.join(dst, os.path.basename(left).replace('L-', ''))
        img = concate_left_right_imgs(
            # left, right, cat_dir, result2, lp=607, rp=319
            left, right, cat_dir, result2, lp=L, rp=R
            )
